/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming.state;



import com.bff.gaia.unified.runners.core.StateInternals;

import com.bff.gaia.unified.runners.core.StateNamespace;

import com.bff.gaia.unified.runners.core.StateTag;

import com.bff.gaia.unified.runners.gaia.translation.types.CoderTypeInformation;

import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.coders.ListCoder;

import com.bff.gaia.unified.sdk.coders.MapCoder;

import com.bff.gaia.unified.sdk.coders.StringUtf8Coder;

import com.bff.gaia.unified.sdk.state.BagState;

import com.bff.gaia.unified.sdk.state.CombiningState;

import com.bff.gaia.unified.sdk.state.MapState;

import com.bff.gaia.unified.sdk.state.ReadableState;

import com.bff.gaia.unified.sdk.state.SetState;

import com.bff.gaia.unified.sdk.state.State;

import com.bff.gaia.unified.sdk.state.StateContext;

import com.bff.gaia.unified.sdk.state.ValueState;

import com.bff.gaia.unified.sdk.state.WatermarkHoldState;

import com.bff.gaia.unified.sdk.transforms.Combine;

import com.bff.gaia.unified.sdk.transforms.CombineWithContext;

import com.bff.gaia.unified.sdk.transforms.windowing.TimestampCombiner;

import com.bff.gaia.unified.sdk.util.CombineContextFactory;

import com.bff.gaia.api.common.ExecutionConfig;

import com.bff.gaia.api.common.state.ListState;

import com.bff.gaia.api.common.state.ListStateDescriptor;

import com.bff.gaia.api.common.state.OperatorStateStore;

import com.bff.gaia.runtime.state.OperatorStateBackend;



import javax.annotation.Nullable;

import java.util.*;



/**

 * {@link StateInternals} that uses a Gaia {@link OperatorStateBackend} to manage the broadcast

 * state. The state is the same on all parallel instances of the operator. So we just need store

 * state of operator-0 in OperatorStateBackend.

 *

 * <p>Note: Ignore index of key. Mainly for SideInputs.

 */

public class GaiaBroadcastStateInternals<K> implements StateInternals {



  private int indexInSubtaskGroup;

  private final OperatorStateBackend stateBackend;

  // stateName -> <namespace, state>

  private Map<String, Map<String, ?>> stateForNonZeroOperator;



  public GaiaBroadcastStateInternals(int indexInSubtaskGroup, OperatorStateBackend stateBackend) {

    this.stateBackend = stateBackend;

    this.indexInSubtaskGroup = indexInSubtaskGroup;

    if (indexInSubtaskGroup != 0) {

      stateForNonZeroOperator = new HashMap<>();

    }

  }



  @Override

  @Nullable

  public K getKey() {

    return null;

  }



  @Override

  public <T extends State> T state(

      final StateNamespace namespace, StateTag<T> address, final StateContext<?> context) {



    return address.bind(

        new StateTag.StateBinder() {



          @Override

          public <T2> ValueState<T2> bindValue(StateTag<ValueState<T2>> address, Coder<T2> coder) {



            return new GaiaBroadcastValueState<>(stateBackend, address, namespace, coder);

          }



          @Override

          public <T2> BagState<T2> bindBag(StateTag<BagState<T2>> address, Coder<T2> elemCoder) {



            return new GaiaBroadcastBagState<>(stateBackend, address, namespace, elemCoder);

          }



          @Override

          public <T2> SetState<T2> bindSet(StateTag<SetState<T2>> address, Coder<T2> elemCoder) {

            throw new UnsupportedOperationException(

                String.format("%s is not supported", SetState.class.getSimpleName()));

          }



          @Override

          public <KeyT, ValueT> MapState<KeyT, ValueT> bindMap(

              StateTag<MapState<KeyT, ValueT>> spec,

              Coder<KeyT> mapKeyCoder,

              Coder<ValueT> mapValueCoder) {

            throw new UnsupportedOperationException(

                String.format("%s is not supported", MapState.class.getSimpleName()));

          }



          @Override

          public <InputT, AccumT, OutputT>

		  CombiningState<InputT, AccumT, OutputT> bindCombiningValue(

                  StateTag<CombiningState<InputT, AccumT, OutputT>> address,

                  Coder<AccumT> accumCoder,

                  Combine.CombineFn<InputT, AccumT, OutputT> combineFn) {



            return new GaiaCombiningState<>(

                stateBackend, address, combineFn, namespace, accumCoder);

          }



          @Override

          public <InputT, AccumT, OutputT>

              CombiningState<InputT, AccumT, OutputT> bindCombiningValueWithContext(

                  StateTag<CombiningState<InputT, AccumT, OutputT>> address,

                  Coder<AccumT> accumCoder,

                  CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn) {

            return new GaiaCombiningStateWithContext<>(

                stateBackend,

                address,

                combineFn,

                namespace,

                accumCoder,

                GaiaBroadcastStateInternals.this,

                CombineContextFactory.createFromStateContext(context));

          }



          @Override

          public WatermarkHoldState bindWatermark(

              StateTag<WatermarkHoldState> address, TimestampCombiner timestampCombiner) {

            throw new UnsupportedOperationException(

                String.format("%s is not supported", WatermarkHoldState.class.getSimpleName()));

          }

        });

  }



  /**

   * 1. The way we would use it is to only checkpoint anything from the operator with subtask index

   * 0 because we assume that the state is the same on all parallel instances of the operator.

   *

   * <p>2. Use map to support namespace.

   */

  private abstract class AbstractBroadcastState<T> {



    private String name;

    private final StateNamespace namespace;

    private final ListStateDescriptor<Map<String, T>> gaiaStateDescriptor;

    private final OperatorStateStore gaiaStateBackend;



    AbstractBroadcastState(

        OperatorStateBackend gaiaStateBackend,

        String name,

        StateNamespace namespace,

        Coder<T> coder) {

      this.name = name;



      this.namespace = namespace;

      this.gaiaStateBackend = gaiaStateBackend;



      CoderTypeInformation<Map<String, T>> typeInfo =

          new CoderTypeInformation<>(MapCoder.of(StringUtf8Coder.of(), coder));



      gaiaStateDescriptor =

          new ListStateDescriptor<>(name, typeInfo.createSerializer(new ExecutionConfig()));

    }



    /** Get map(namespce->T) from index 0. */

    Map<String, T> getMap() throws Exception {

      if (indexInSubtaskGroup == 0) {

        return getMapFromBroadcastState();

      } else {

        Map<String, T> result = (Map<String, T>) stateForNonZeroOperator.get(name);

        // maybe restore from BroadcastState of Operator-0

        if (result == null) {

          result = getMapFromBroadcastState();

          if (result != null) {

            stateForNonZeroOperator.put(name, result);

            // we don't need it anymore, must clear it.

            gaiaStateBackend.getUnionListState(gaiaStateDescriptor).clear();

          }

        }

        return result;

      }

    }



    Map<String, T> getMapFromBroadcastState() throws Exception {

      ListState<Map<String, T>> state = gaiaStateBackend.getUnionListState(gaiaStateDescriptor);

      Iterable<Map<String, T>> iterable = state.get();

      Map<String, T> ret = null;

      if (iterable != null) {

        // just use index 0

        Iterator<Map<String, T>> iterator = iterable.iterator();

        if (iterator.hasNext()) {

          ret = iterator.next();

        }

      }

      return ret;

    }



    /** Update map(namespce->T) from index 0. */

    void updateMap(Map<String, T> map) throws Exception {

      if (indexInSubtaskGroup == 0) {

        ListState<Map<String, T>> state = gaiaStateBackend.getUnionListState(gaiaStateDescriptor);

        state.clear();

        if (map.size() > 0) {

          state.add(map);

        }

      } else {

        if (map.isEmpty()) {

          stateForNonZeroOperator.remove(name);

          // updateMap is always behind getMap,

          // getMap will clear map in BroadcastOperatorState,

          // we don't need clear here.

        } else {

          stateForNonZeroOperator.put(name, map);

        }

      }

    }



    void writeInternal(T input) {

      try {

        Map<String, T> map = getMap();

        if (map == null) {

          map = new HashMap<>();

        }

        map.put(namespace.stringKey(), input);

        updateMap(map);

      } catch (Exception e) {

        throw new RuntimeException("Error updating state.", e);

      }

    }



    T readInternal() {

      try {

        Map<String, T> map = getMap();

        if (map == null) {

          return null;

        } else {

          return map.get(namespace.stringKey());

        }

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    void clearInternal() {

      try {

        Map<String, T> map = getMap();

        if (map != null) {

          map.remove(namespace.stringKey());

          updateMap(map);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error clearing state.", e);

      }

    }

  }



  private class GaiaBroadcastValueState<T> extends AbstractBroadcastState<T>

      implements ValueState<T> {



    private final StateNamespace namespace;

    private final StateTag<ValueState<T>> address;



    GaiaBroadcastValueState(

        OperatorStateBackend gaiaStateBackend,

        StateTag<ValueState<T>> address,

        StateNamespace namespace,

        Coder<T> coder) {

      super(gaiaStateBackend, address.getId(), namespace, coder);



      this.namespace = namespace;

      this.address = address;

    }



    @Override

    public void write(T input) {

      writeInternal(input);

    }



    @Override

    public ValueState<T> readLater() {

      return this;

    }



    @Override

    public T read() {

      return readInternal();

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaBroadcastValueState<?> that = (GaiaBroadcastValueState<?>) o;



      return namespace.equals(that.namespace) && address.equals(that.address);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + address.hashCode();

      return result;

    }



    @Override

    public void clear() {

      clearInternal();

    }

  }



  private class GaiaBroadcastBagState<T> extends AbstractBroadcastState<List<T>>

      implements BagState<T> {



    private final StateNamespace namespace;

    private final StateTag<BagState<T>> address;



    GaiaBroadcastBagState(

        OperatorStateBackend gaiaStateBackend,

        StateTag<BagState<T>> address,

        StateNamespace namespace,

        Coder<T> coder) {

      super(gaiaStateBackend, address.getId(), namespace, ListCoder.of(coder));



      this.namespace = namespace;

      this.address = address;

    }



    @Override

    public void add(T input) {

      List<T> list = readInternal();

      if (list == null) {

        list = new ArrayList<>();

      }

      list.add(input);

      writeInternal(list);

    }



    @Override

    public BagState<T> readLater() {

      return this;

    }



    @Override

    public Iterable<T> read() {

      List<T> result = readInternal();

      return result != null ? result : Collections.emptyList();

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            List<T> result = readInternal();

            return result == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      clearInternal();

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaBroadcastBagState<?> that = (GaiaBroadcastBagState<?>) o;



      return namespace.equals(that.namespace) && address.equals(that.address);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + address.hashCode();

      return result;

    }

  }



  private class GaiaCombiningState<InputT, AccumT, OutputT> extends AbstractBroadcastState<AccumT>

      implements CombiningState<InputT, AccumT, OutputT> {



    private final StateNamespace namespace;

    private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;

    private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn;



    GaiaCombiningState(

        OperatorStateBackend gaiaStateBackend,

        StateTag<CombiningState<InputT, AccumT, OutputT>> address,

        Combine.CombineFn<InputT, AccumT, OutputT> combineFn,

        StateNamespace namespace,

        Coder<AccumT> accumCoder) {

      super(gaiaStateBackend, address.getId(), namespace, accumCoder);



      this.namespace = namespace;

      this.address = address;

      this.combineFn = combineFn;

    }



    @Override

    public CombiningState<InputT, AccumT, OutputT> readLater() {

      return this;

    }



    @Override

    public void add(InputT value) {

      AccumT current = readInternal();

      if (current == null) {

        current = combineFn.createAccumulator();

      }

      current = combineFn.addInput(current, value);

      writeInternal(current);

    }



    @Override

    public void addAccum(AccumT accum) {

      AccumT current = readInternal();



      if (current == null) {

        writeInternal(accum);

      } else {

        current = combineFn.mergeAccumulators(Arrays.asList(current, accum));

        writeInternal(current);

      }

    }



    @Override

    public AccumT getAccum() {

      AccumT accum = readInternal();

      return accum != null ? accum : combineFn.createAccumulator();

    }



    @Override

    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {

      return combineFn.mergeAccumulators(accumulators);

    }



    @Override

    public OutputT read() {

      AccumT accum = readInternal();

      if (accum != null) {

        return combineFn.extractOutput(accum);

      } else {

        return combineFn.extractOutput(combineFn.createAccumulator());

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return readInternal() == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      clearInternal();

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaCombiningState<?, ?, ?> that = (GaiaCombiningState<?, ?, ?>) o;



      return namespace.equals(that.namespace) && address.equals(that.address);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + address.hashCode();

      return result;

    }

  }



  private class GaiaKeyedCombiningState<K2, InputT, AccumT, OutputT>

      extends AbstractBroadcastState<AccumT> implements CombiningState<InputT, AccumT, OutputT> {



    private final StateNamespace namespace;

    private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;

    private final Combine.CombineFn<InputT, AccumT, OutputT> combineFn;

    private final GaiaBroadcastStateInternals<K2> gaiaStateInternals;



    GaiaKeyedCombiningState(

        OperatorStateBackend gaiaStateBackend,

        StateTag<CombiningState<InputT, AccumT, OutputT>> address,

        Combine.CombineFn<InputT, AccumT, OutputT> combineFn,

        StateNamespace namespace,

        Coder<AccumT> accumCoder,

        GaiaBroadcastStateInternals<K2> gaiaStateInternals) {

      super(gaiaStateBackend, address.getId(), namespace, accumCoder);



      this.namespace = namespace;

      this.address = address;

      this.combineFn = combineFn;

      this.gaiaStateInternals = gaiaStateInternals;

    }



    @Override

    public CombiningState<InputT, AccumT, OutputT> readLater() {

      return this;

    }



    @Override

    public void add(InputT value) {

      try {

        AccumT current = readInternal();

        if (current == null) {

          current = combineFn.createAccumulator();

        }

        current = combineFn.addInput(current, value);

        writeInternal(current);

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public void addAccum(AccumT accum) {

      try {

        AccumT current = readInternal();

        if (current == null) {

          writeInternal(accum);

        } else {

          current = combineFn.mergeAccumulators(Arrays.asList(current, accum));

          writeInternal(current);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public AccumT getAccum() {

      try {

        AccumT accum = readInternal();

        return accum != null ? accum : combineFn.createAccumulator();

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {

      return combineFn.mergeAccumulators(accumulators);

    }



    @Override

    public OutputT read() {

      try {

        AccumT accum = readInternal();

        if (accum != null) {

          return combineFn.extractOutput(accum);

        } else {

          return combineFn.extractOutput(combineFn.createAccumulator());

        }

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return readInternal() == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      clearInternal();

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaKeyedCombiningState<?, ?, ?, ?> that = (GaiaKeyedCombiningState<?, ?, ?, ?>) o;



      return namespace.equals(that.namespace) && address.equals(that.address);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + address.hashCode();

      return result;

    }

  }



  private class GaiaCombiningStateWithContext<K2, InputT, AccumT, OutputT>

      extends AbstractBroadcastState<AccumT> implements CombiningState<InputT, AccumT, OutputT> {



    private final StateNamespace namespace;

    private final StateTag<CombiningState<InputT, AccumT, OutputT>> address;

    private final CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn;

    private final GaiaBroadcastStateInternals<K2> gaiaStateInternals;

    private final CombineWithContext.Context context;



    GaiaCombiningStateWithContext(

        OperatorStateBackend gaiaStateBackend,

        StateTag<CombiningState<InputT, AccumT, OutputT>> address,

        CombineWithContext.CombineFnWithContext<InputT, AccumT, OutputT> combineFn,

        StateNamespace namespace,

        Coder<AccumT> accumCoder,

        GaiaBroadcastStateInternals<K2> gaiaStateInternals,

        CombineWithContext.Context context) {

      super(gaiaStateBackend, address.getId(), namespace, accumCoder);



      this.namespace = namespace;

      this.address = address;

      this.combineFn = combineFn;

      this.gaiaStateInternals = gaiaStateInternals;

      this.context = context;

    }



    @Override

    public CombiningState<InputT, AccumT, OutputT> readLater() {

      return this;

    }



    @Override

    public void add(InputT value) {

      try {

        AccumT current = readInternal();

        if (current == null) {

          current = combineFn.createAccumulator(context);

        }

        current = combineFn.addInput(current, value, context);

        writeInternal(current);

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public void addAccum(AccumT accum) {

      try {



        AccumT current = readInternal();

        if (current == null) {

          writeInternal(accum);

        } else {

          current = combineFn.mergeAccumulators(Arrays.asList(current, accum), context);

          writeInternal(current);

        }

      } catch (Exception e) {

        throw new RuntimeException("Error adding to state.", e);

      }

    }



    @Override

    public AccumT getAccum() {

      try {

        AccumT accum = readInternal();

        return accum != null ? accum : combineFn.createAccumulator(context);

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public AccumT mergeAccumulators(Iterable<AccumT> accumulators) {

      return combineFn.mergeAccumulators(accumulators, context);

    }



    @Override

    public OutputT read() {

      try {

        AccumT accum = readInternal();

        if (accum == null) {

          accum = combineFn.createAccumulator(context);

        }

        return combineFn.extractOutput(accum, context);

      } catch (Exception e) {

        throw new RuntimeException("Error reading state.", e);

      }

    }



    @Override

    public ReadableState<Boolean> isEmpty() {

      return new ReadableState<Boolean>() {

        @Override

        public Boolean read() {

          try {

            return readInternal() == null;

          } catch (Exception e) {

            throw new RuntimeException("Error reading state.", e);

          }

        }



        @Override

        public ReadableState<Boolean> readLater() {

          return this;

        }

      };

    }



    @Override

    public void clear() {

      clearInternal();

    }



    @Override

    public boolean equals(Object o) {

      if (this == o) {

        return true;

      }

      if (o == null || getClass() != o.getClass()) {

        return false;

      }



      GaiaCombiningStateWithContext<?, ?, ?, ?> that =

          (GaiaCombiningStateWithContext<?, ?, ?, ?>) o;



      return namespace.equals(that.namespace) && address.equals(that.address);

    }



    @Override

    public int hashCode() {

      int result = namespace.hashCode();

      result = 31 * result + address.hashCode();

      return result;

    }

  }

}